#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@author: liang kang
@contact: gangkanli1219@gmail.com
@time: 2018/1/26 16:10
@desc: 
"""
import os

import numpy as np
import tensorflow as tf
from skimage import transform

from .utils.eval import check_image_symbol
from .utils.imaging import convert_density_map
from .utils.imaging import preprocess_image
from .utils.imaging import rotate_image
from .utils.json import data_to_json_high


def predict_image(image, model, sub_shape=(320, 320), input_shape=(256, 256), scale=1.0, shape=(1280, 1280)):
    """
    预测图像中的人群密度和数量

    Parameters
    ----------
    image: 待预测的图像，可以为图像路径或者numpy的array
    model:
    sub_shape: 大图切割的子图的大小
    input_shape: 输入网络的图像大小
    scale: float, resize_shape / input_shape, 取值(0,1]. 决定子图在缩小时相当于input_shape的大小
    shape: 预定义的鱼眼图像大小

    Returns
    -------
    mask: density map
    """
    mask_factor = np.load(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'data', 'high', 'mask.npy'))

    resize_shape = (int(round(input_shape[0] * scale)), int(round(input_shape[1] * scale)))

    config = tf.ConfigProto(device_count={"CPU": 4, "GPU": 1})
    config.gpu_options.allow_growth = True
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
    with tf.Session(config=config) as sess:

        meta_graph_def = tf.saved_model.loader.load(sess, ['frozen_model'], model)
        signature = meta_graph_def.signature_def
        image_tensor_name = signature['predict_images'].inputs['images'].name
        pre_density_map_name = signature['predict_images'].outputs['density_map'].name
        image_tensor = sess.graph.get_tensor_by_name(image_tensor_name)
        density_map_name = sess.graph.get_tensor_by_name(pre_density_map_name)

        mask = np.zeros(shape).astype(np.float32)
        for angle in range(0, 180, 30):
            offset, img_buf = rotate_image(angle, image)
            img1 = np.expand_dims(preprocess_image(img_buf[40:360, 480:800, :],
                                                   shape=input_shape, resize_shape=resize_shape), axis=0)
            img2 = np.expand_dims(preprocess_image(img_buf[920:1240, 480:800, :],
                                                   shape=input_shape, resize_shape=resize_shape), axis=0)
            density_map1 = sess.run([density_map_name], feed_dict={image_tensor: img1})
            density_map2 = sess.run([density_map_name], feed_dict={image_tensor: img2})
            buffer = np.zeros(shape)
            buffer[40:360, 480:800] = convert_density_map(np.squeeze(density_map1),
                                                          shape=sub_shape, resize_shape=resize_shape)
            buffer[920:1240, 480:800] = convert_density_map(np.squeeze(density_map2),
                                                            shape=sub_shape, resize_shape=resize_shape)
            if 0 != angle:
                buffer = transform.rotate(buffer, -angle, resize=True, preserve_range=True)
                buffer = buffer[offset[0]:(offset[0] + shape[0]), offset[1]:(offset[1] + shape[1])]
            mask += buffer

        small_shape = (720, 720)
        small_img = image[280:1000, 280:1000, :]
        for angle in range(0, 360, 72):
            offset, small_img_buf = rotate_image(angle, small_img)
            img_patch = np.expand_dims(preprocess_image(small_img_buf[80:400, 240:560, :],
                                                        shape=input_shape, resize_shape=resize_shape), axis=0)
            density_map = sess.run([density_map_name], feed_dict={image_tensor: img_patch})
            buffer = np.zeros(small_shape)
            buffer[80:400, 240:560] = convert_density_map(np.squeeze(density_map),
                                                          shape=sub_shape, resize_shape=resize_shape)
            if 0 != angle:
                buffer = transform.rotate(buffer, -angle, resize=True, preserve_range=True)
                buffer = buffer[offset[0]:(offset[0] + small_shape[0]), offset[1]:(offset[1] + small_shape[1])]
            mask[280:1000, 280:1000] += buffer

        return mask * mask_factor


def predict_high_density_image(image, rois=None, high=7.0, store_id=0, types=None,
                               shape=(1280, 1280)):
    """
    预测高密度图像, 如果rois为None，返回整张图像的人数，否则返回一个字典
    这个字典的value就是rois中对应key的roi的预测人数

    Parameters
    ----------
    image: 输入图像
    rois: roi区域
    high: 摄像头高度，[4, 7]
    store_id
    types
    shape: 预定义的鱼眼图像大小

    Returns
    -------

    """
    model = os.path.join(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'data', 'high', '73205'))
    assert 4 <= high, 'the height is too low !'
    assert 7 >= high, 'the height is too high !'
    scale = 0.1 * high + 0.3
    img = check_image_symbol(image)
    density_map = predict_image(img, model, scale=scale, shape=shape)
    return data_to_json_high(image, density_map, rois, store_id, types)


